package cn.iocoder.yudao.framework.datasource.config;

import cn.iocoder.yudao.framework.datasource.provider.AbstractDataSourceProvider;
import cn.iocoder.yudao.framework.datasource.provider.DynamicDataSourceProvider;
import com.baomidou.dynamic.datasource.ds.AbstractRoutingDataSource;
import com.baomidou.dynamic.datasource.ds.GroupDataSource;
import com.baomidou.dynamic.datasource.ds.ItemDataSource;
import com.baomidou.dynamic.datasource.exception.CannotFindDataSourceException;
import com.baomidou.dynamic.datasource.spring.boot.autoconfigure.DataSourceProperty;
import com.baomidou.dynamic.datasource.strategy.DynamicDataSourceStrategy;
import com.baomidou.dynamic.datasource.strategy.LoadBalanceDynamicDataSourceStrategy;
import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.util.ReflectionUtils;
import org.springframework.util.StringUtils;

import javax.sql.DataSource;
import java.lang.reflect.Method;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * AbstractRoutingDataSource从版本3.5.1继承
 *
 * @className: DynamicRoutingDataSource
 * @author: tangkh
 * @date: 2023-09-19 16:55
 */
public class DynamicRoutingDataSource extends AbstractRoutingDataSource implements InitializingBean, DisposableBean {
    private static final Logger log = LoggerFactory.getLogger(DynamicRoutingDataSource.class);
    private static final String UNDERLINE = "_";
    private final Map<String, DataSource> dataSourceMap = new ConcurrentHashMap();
    private final Map<String, GroupDataSource> groupDataSources = new ConcurrentHashMap();
    @Autowired
    private DynamicDataSourceProvider provider;
    private Class<? extends DynamicDataSourceStrategy> strategy = LoadBalanceDynamicDataSourceStrategy.class;
    private String primary = "master";
    private Boolean strict = false;
    private Boolean p6spy = false;
    private Boolean seata = false;

    public DynamicRoutingDataSource() {
    }

    @Override
    public DataSource determineDataSource() {
        String dsKey = DynamicDataSourceContextHolder.peek();
        return this.getDataSource(dsKey);
    }

    @Override
    protected String getPrimary() {
        return null;
    }

    private DataSource determinePrimaryDataSource() {
        log.debug("dynamic-datasource switch to the primary datasource");
        DataSource dataSource = this.dataSourceMap.get(this.primary);
        if (dataSource != null) {
            return dataSource;
        } else {
            GroupDataSource groupDataSource = this.groupDataSources.get(this.primary);
            if (groupDataSource != null) {
                return groupDataSource.determineDataSource();
            } else {
                throw new CannotFindDataSourceException("dynamic-datasource can not find primary datasource");
            }
        }
    }

    public Map<String, DataSource> getDataSources() {
        return this.dataSourceMap;
    }

    public Map<String, GroupDataSource> getGroupDataSources() {
        return this.groupDataSources;
    }

    public DataSource getDataSource(String ds) {
        if (StringUtils.isEmpty(ds)) {
            return this.determinePrimaryDataSource();
        } else if (this.dataSourceMap.containsKey(ds)) {
            log.debug("dynamic-datasource switch to the datasource named [{}]", ds);
            return this.dataSourceMap.get(ds);
        } else if (this.strict) {
            throw new CannotFindDataSourceException("dynamic-datasource could not find a datasource named:" + ds);
        } else {
            return this.determinePrimaryDataSource();
        }
    }

    public synchronized void addDataSource(String ds, DataSource dataSource) {
        DataSource oldDataSource = this.dataSourceMap.put(ds, dataSource);
        this.addGroupDataSource(ds, dataSource);
        if (oldDataSource != null) {
            this.closeDataSource(ds, oldDataSource);
        }

        log.info("dynamic-datasource - add a datasource named [{}] success", ds);
    }

    private void addGroupDataSource(String ds, DataSource dataSource) {
        if (ds.contains("_")) {
            String group = ds.split("_")[0];
            GroupDataSource groupDataSource = this.groupDataSources.get(group);
            if (groupDataSource == null) {
                try {
                    groupDataSource = new GroupDataSource(group, this.strategy.getDeclaredConstructor().newInstance());
                    this.groupDataSources.put(group, groupDataSource);
                } catch (Exception var6) {
                    throw new RuntimeException("dynamic-datasource - add the datasource named " + ds + " error", var6);
                }
            }

            groupDataSource.addDatasource(ds, dataSource);
        }

    }

    public synchronized void removeDataSource(String ds) {
        if (!StringUtils.hasText(ds)) {
            throw new RuntimeException("remove parameter could not be empty");
        } else if (this.primary.equals(ds)) {
            throw new RuntimeException("could not remove primary datasource");
        } else {
            if (this.dataSourceMap.containsKey(ds)) {
                DataSource dataSource = this.dataSourceMap.remove(ds);
                this.closeDataSource(ds, dataSource);
                if (ds.contains("_")) {
                    String group = ds.split("_")[0];
                    if (this.groupDataSources.containsKey(group)) {
                        DataSource oldDataSource = this.groupDataSources.get(group).removeDatasource(ds);
                        if (oldDataSource == null) {
                            log.warn("fail for remove datasource from group. dataSource: {} ,group: {}", ds, group);
                        }
                    }
                }

                log.info("dynamic-datasource - remove the database named [{}] success", ds);
            } else {
                log.warn("dynamic-datasource - could not find a database named [{}]", ds);
            }

        }
    }

    @Override
    public void destroy() {
        log.info("dynamic-datasource start closing ....");
        Iterator var1 = this.dataSourceMap.entrySet().iterator();

        while (var1.hasNext()) {
            Map.Entry<String, DataSource> item = (Map.Entry) var1.next();
            this.closeDataSource(item.getKey(), item.getValue());
        }

        log.info("dynamic-datasource all closed success,bye");
    }

    @Override
    public void afterPropertiesSet() {
        Map<String, DataSource> dataSources = new HashMap(16);
        Map<String, DataSourceProperty> dataSourcePropertyMap = new ConcurrentHashMap<>();
        Map dspMap = this.provider.dataSourceProperties();
        if (dspMap != null) {
            dataSourcePropertyMap.putAll(dspMap);
        }

        Map<String, DataSource> loadDataSources = this.provider.loadDataSources();
        if (loadDataSources == null) {
            log.error("没有装载到任何数据源配置，请检查JDBC配置或者YML");
        } else {
            dataSources.putAll(loadDataSources);
            Iterator var5 = dataSources.entrySet().iterator();

            while (var5.hasNext()) {
                Map.Entry<String, DataSource> dsItem = (Map.Entry) var5.next();
                this.addDataSource(dsItem.getKey(), dsItem.getValue());
            }

            if (this.groupDataSources.containsKey(this.primary)) {
                log.info("dynamic-datasource initial loaded [{}] datasource,primary group datasource named [{}]", dataSources.size(), this.primary);
            } else if (this.dataSourceMap.containsKey(this.primary)) {
                log.info("dynamic-datasource initial loaded [{}] datasource,primary datasource named [{}]", dataSources.size(), this.primary);
            } else {
                log.warn("dynamic-datasource initial loaded [{}] datasource,Please add your primary datasource or check your configuration", dataSources.size());
            }

            ScheduledThreadPoolExecutor scheduledThreadPoolExecutor = new ScheduledThreadPoolExecutor(1);
            scheduledThreadPoolExecutor.scheduleAtFixedRate(() -> {
                Map<String, DataSourceProperty> newDataSources = new HashMap();
                Map newDsmap = this.provider.dataSourceProperties();
                if (newDsmap != null) {
                    newDataSources.putAll(newDsmap);
                }

                AbstractDataSourceProvider abstractDataSourceProvider = (AbstractDataSourceProvider) this.provider;
                dataSourcePropertyMap.forEach((k, v) -> {
                    DataSourceProperty dataSourceProperty = dataSourcePropertyMap.get(k);
                    DataSourceProperty newDataSourceProperty = newDataSources.get(k);
                    if (newDataSourceProperty == null) {
                        log.info("删除数据源:{}", k);
                        this.removeDataSource(k);
                        dataSourcePropertyMap.remove(k);
                    } else {
                        String var10000 = newDataSourceProperty.getPoolName();
                        String newKey = var10000 + newDataSourceProperty.getUrl() + newDataSourceProperty.getUsername() + newDataSourceProperty.getPassword();
                        var10000 = dataSourceProperty.getPoolName();
                        String dataSourceKey = var10000 + dataSourceProperty.getUrl() + dataSourceProperty.getUsername() + dataSourceProperty.getPassword();
                        if (!newKey.equals(dataSourceKey)) {
                            Map<String, DataSourceProperty> dataSourcePropertiesMap = new HashMap();
                            dataSourcePropertiesMap.put(k, newDataSourceProperty);
                            Map<String, DataSource> providerDataSourceMap = abstractDataSourceProvider.createDataSourceMap(dataSourcePropertiesMap);
                            log.info("修改数据源:{}", k);
                            this.addDataSource(k, providerDataSourceMap.get(k));
                            dataSourcePropertyMap.put(k, newDataSourceProperty);
                        }
                    }

                });
                newDataSources.forEach((k, v) -> {
                    DataSourceProperty dataSourceProperty = dataSourcePropertyMap.get(k);
                    DataSourceProperty newDataSourceProperty = newDataSources.get(k);
                    if (dataSourceProperty == null) {
                        Map<String, DataSourceProperty> dataSourcePropertiesMap = new HashMap<>();
                        dataSourcePropertiesMap.put(k, newDataSourceProperty);
                        Map<String, DataSource> providerDataSourceMap = abstractDataSourceProvider.createDataSourceMap(dataSourcePropertiesMap);
                        log.info("新增数据源:{}", k);
                        this.addDataSource(k, providerDataSourceMap.get(k));
                        dataSourcePropertyMap.put(k, newDataSourceProperty);
                    }

                });
            }, 10L, 30L, TimeUnit.SECONDS);
        }
    }

    private void closeDataSource(String ds, DataSource dataSource) {
        try {
            if (dataSource instanceof ItemDataSource) {
                ((ItemDataSource) dataSource).close();
            } else {
                Method closeMethod = ReflectionUtils.findMethod(dataSource.getClass(), "close");
                if (closeMethod != null) {
                    closeMethod.invoke(dataSource);
                }
            }
        } catch (Exception var4) {
            log.warn("dynamic-datasource closed datasource named [{}] failed", ds, var4);
        }

    }

    public void setStrategy(Class<? extends DynamicDataSourceStrategy> strategy) {
        this.strategy = strategy;
    }

    public void setPrimary(String primary) {
        this.primary = primary;
    }

    public void setStrict(Boolean strict) {
        this.strict = strict;
    }

    public void setP6spy(Boolean p6spy) {
        this.p6spy = p6spy;
    }

    public void setSeata(Boolean seata) {
        this.seata = seata;
    }
}
